import cv2
import numpy as np
import asyncio
import threading
import time
from queue import Queue
from fastapi import FastAPI, WebSocket
from fastapi.responses import HTMLResponse
from fastapi.staticfiles import StaticFiles
from fastapi.middleware.cors import CORSMiddleware
from starlette.endpoints import WebSocketEndpoint
from starlette.websockets import WebSocketDisconnect
from go2_webrtc_driver.webrtc_driver import Go2WebRTCConnection, WebRTCConnectionMethod
from aiortc import MediaStreamTrack
from ultralytics import YOLO
import uvicorn
import base64

app = FastAPI()
app.add_middleware(CORSMiddleware, allow_origins=["*"], allow_methods=["*"], allow_headers=["*"])
app.mount("/static", StaticFiles(directory="static"), name="static")

model = YOLO("yolo11n.pt")
print("✅ YOLOv11n 模型加载成功，支持类别：", model.names)

frame_queue = Queue()
conn = Go2WebRTCConnection(WebRTCConnectionMethod.LocalSTA, ip="192.168.123.161")

@app.get("/")
async def get():
    with open("static/index.html") as f:
        return HTMLResponse(f.read())

@app.websocket("/ws")
class VideoStreamer(WebSocketEndpoint):
    encoding = "text"

    async def on_connect(self, websocket: WebSocket):
        await websocket.accept()
        self.mode = "ai"  # 默认开启AI识别

    async def on_receive(self, websocket: WebSocket, data):
        if data in ["raw", "ai"]:
            self.mode = data

    async def on_disconnect(self, websocket: WebSocket, close_code: int):
        pass

    async def on_receive_bytes(self, websocket: WebSocket, data: bytes):
        pass

    async def dispatch(self):
        try:
            frame_counter = 0
            while True:
                if not frame_queue.empty():
                    img = frame_queue.get()

                    # 原始或识别结果图像
                    if self.mode == "ai" and frame_counter % 9 == 0:
                        results = model(img, verbose=False)[0]
                        img = results.plot()

                    frame_counter += 1

                    # 编码为JPEG并推送给前端
                    _, jpeg = cv2.imencode('.jpg', img)
                    jpg_as_text = base64.b64encode(jpeg.tobytes()).decode('utf-8')
                    await websocket.send_text(jpg_as_text)
                await asyncio.sleep(0.03)
        except WebSocketDisconnect:
            return

async def recv_camera_stream(track: MediaStreamTrack):
    while True:
        frame = await track.recv()
        img = frame.to_ndarray(format="bgr24")
        frame_queue.put(img)

def run_webrtc():
    asyncio.set_event_loop(asyncio.new_event_loop())
    loop = asyncio.get_event_loop()

    async def setup():
        await conn.connect()
        conn.video.switchVideoChannel(True)
        conn.video.add_track_callback(recv_camera_stream)

    loop.run_until_complete(setup())
    loop.run_forever()

threading.Thread(target=run_webrtc, daemon=True).start()

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=8000)
